/* Mother function for HZ dispersal model.
Code based very loosely on Ben and Jeremy's LD code
 
 Space is one dimensional
*/

#include <R.h>
#include <Rinternals.h>
#include <Rmath.h>
#include <math.h>
#include <Rdefines.h>


//------------------------------------------------------------------------------
typedef struct{
	double x_coord;
	int alive;
	double HI;
	double disp_B; //Breeding value for dispersal
	double C1; //Choosiness -- because eventually we want individual variability in this trait
	double C2;
} ind;

//Global variables (they're back!!)
ind *new_ind = NULL; //next generation of individuals
ind *cur_ind = NULL;
int **density=NULL;
int ***indices=NULL;


//Function Declarations------------------------------------------

void initialise (int n, int initial_col, int xDim);
void grid_update(int n, int xDim);
double choosiness (double diffs, double C, int a, int k);
int mate_selection(int ii, int xx, double C);
double hybrid_fitness(double HI, double s, double beta);
int reproduce(double maxcost, double beta, double mut_rate, int xDim, int popsize, double lambda, double K, double Vadd_disp);
void mutations(ind * individuals, int id, double mut_rate, int xDim);
void disperse(int xDim, double mort, int popsize);
void check_inds (ind * individuals, int spX, int n);
void check_grid(int **gd, int spX, int n);
void check_indices(int ** gd, int *** in, int spX);
int first_empty(int * tt, int len);
void new_to_curr(int n);
void output_individuals(int n, char mname[10], SEXP env);
//double c_gauss_rng(double sd);
int disp_mort(double dist);
void disp_summ(char vname[10], int popsize, int xDim, SEXP env);
//---------------------------------------------------------------


//Function Definitions

void initialise (int n, int initial_col, int xDim){
  //initialize the density grid
	density =  Calloc(xDim, int *); //allocate memory for x dimension
  	
  	// initialise the index grid
	indices = Calloc(xDim, int*); for (int xx=0; xx<xDim; xx++) { 
        indices[xx] = Calloc(density[xx],int); //allocate 2nd dim of this matrix
        for (int zz=0; zz<density[xx]; zz++) { indices[xx][zz] = -99; } //set everything to -99
		}
	}
	//Initialise the individual array
	cur_ind=Calloc(n, ind);
	new_ind=Calloc(1, ind);
  int x;
  //Initialise the first species (n/2 individuals)
  for (int i = 0; i < n/2; i++){
   cur_ind[i].x_coord = Rf_runif(0,initial_col);
   cur_ind[i].alive = 1;
   cur_ind[i].HI = 0.0;
   cur_ind[i].disp_B = Rf_runif(0,xDim/10);
   cur_ind[i].C1 = 0.0; //fixed for now
   cur_ind[i].C2 = 0.0;
  }
  //Initialise the second species (n/2 cur_ind)
  for (int i = (n/2); i < n; i++)
  {
   cur_ind[i].x_coord = Rf_runif(xDim - initial_col,xDim);
   cur_ind[i].alive = 1;
   cur_ind[i].HI = 1.0;
   cur_ind[i].disp_B = Rf_runif(0,xDim/10);
   cur_ind[i].C1 = 0.0; //fixed for now
   cur_ind[i].C2 = 0.0;
  }
  grid_update(n, xDim);

}


//Updates grid with N in each cell and writes individual indices into grid.

void grid_update(int nn, int xDim){
  int x, z;
  for (x=0; x<xDim; x++){density[x]=0;}
  for (int ii=0; ii<nn; ii++){ // Count cur_ind indices each cell
	if (cur_ind[ii].alive==0) {	continue;} //skip deadies
	x = (int)cur_ind[ii].x_coord;
	density[x]++;
  }
	indices=Realloc(indices, xDim, int**);
	for (x=0; x<xDim; x++) { indices[x]=Realloc(indices[x], density[x], int);//reallocate memory to indices
        for (z=0; z<density[x]; z++) { indices[x][z] = -99; } //set everything to -99	
	}
  for (int ii=0; ii<nn; ii++){//Storing the cur_ind indices
	if (cur_ind[ii].alive==0) {	continue;} //skip deadies
	x = (int)cur_ind[ii].x_coord;
    indices[x][first_empty(indices[x],density[x])] = ii;
  }
}


double choosiness (double diffs, double C, int a, int k){
  return (exp(a-C*k*diffs)/(1+exp(a-C*k*diffs)));
}


//ii is the index when going through the individuals array
//C = "chosiness"
//xx & yy are the individual coordinate
//This function is for 1 individual, so it must be used inside an individual loop
int mate_selection(int ii, int xx, double C)
{
 int mi;
 double *HID;
 double *attract;
 double sum_attract = 0.0;
 double *cumPchoose;
 HID = Calloc(density[xx], double); //hybrid index distances
 attract = Calloc(density[xx], double);//mate attractiveness
 cumPchoose = Calloc(density[xx], double); //cumulative choosing probabilities

 if (density[xx]==1) {return(ii);}
 for (int j = 0; j < density[xx]; j++)
 {
  int jj = indices[xx][j] ;

  HID[j]=abs(cur_ind[ii].HI - cur_ind[jj].HI); //absolute difference in HI

  attract[j] = choosiness(HID[j],C,5,100);

  sum_attract += attract[j];
  //printf("HID= %f attract=%f sum=%f pot_mate=%i\n", HID[j], attract[j], sum_attract, jj);
 }
 for (int j = 0; j < density[xx]; j++){
	 if (j==0) {
		 cumPchoose[j] = (attract[j] / sum_attract);
	 }
	 else {
		 cumPchoose[j] = cumPchoose[j-1]+(attract[j] / sum_attract);
	 }

   //printf("cpchoose=%f\n", cumPchoose[j]);
 }
 double rn = runif(0.0,1.0);
 for (mi=0; mi<density[xx];mi++) {
	 if (cumPchoose[mi]>=rn) {break;}
 }

 Free(cumPchoose);
 Free(HID);
 Free(attract);

 return(indices[xx][mi]);
}

// HI = hybrid index from Barton and Gale
double hybrid_fitness(double HI, double s, double beta){
  return(1-s*(4*HI*(1-HI))^beta);
}    


int reproduce(double maxcost,
			   double beta, double mut_rate, int xDim,
			   int popsize, double lambda, double K, double Vadd_disp)
{
 int x;
 int new_n=0;
 int *n_offs = Calloc(popsize, int);
 int nn = 0; //counter for new_inds
 int *MI=Calloc(popsize, int); //mate index
 double E, fitness;

 //Loop through each inds and calculate number of offspring for each individual
 for (int i = 0; i < popsize; i++)
 {
  	{
   		x = (int)cur_ind[i].x_coord;
		//select a mate
		double meanC = (cur_ind[i].C1+cur_ind[i].C2)/2;
		MI[i] = mate_selection(i, x, meanC);
		if (MI[i]==i) {MI[i]=-9; n_offs[i]=0; continue;} //catch instances of no mates
   		//calculate individual fitness (dead individuals have zero fitness)
	   	fitness = cur_ind[i].alive*hybrid_fitness(cur_ind[i].HI, maxcost, beta)//*hybrid_fitness(cur_ind[MI[i]].HI, maxcost, beta);
	   	
   		//hassel commins population growth where b=1 (contest competition)
	   	//generates expected fecundity
   		E = fitness*lambda/(1+(lambda-1)*density[x]/K);
	   	n_offs[i] = Rf_rpois(E);
	   	//printf("noffsp = %i HI= %f fitness= %f, E=%f\n", n_offs[i], cur_ind[i].HI, fitness, E);
   		new_n += n_offs[i];
  	}
 }
 //allocate memory for new individuals
 new_ind = Realloc(new_ind, new_n, ind); 
 for (int i = 0; i < popsize; i++){ 
   if (n_offs[i] > 0)
   {
   		x = (int)cur_ind[i].x_coord;
		//loop through new offspring and initialise them
		for (int j = 0; j < n_offs[i]; j++)
		{
			new_ind[nn].x_coord = cur_ind[i].x_coord;
			new_ind[nn].alive = 1;
			new_ind[nn].HI = (cur_ind[i].HI+cur_ind[MI[i]].HI)/2; // Note: this halves variance in HI each generation!
			new_ind[nn].disp_B = Rf_rnorm((cur_ind[i].disp_B+cur_ind[MI[i]].disp_B)/2, Vadd_disp/2); //Offspring's breeding value
			while (new_ind[nn].disp_B <=0) {new_ind[nn].disp_B = Rf_rnorm((cur_ind[i].disp_B+cur_ind[MI[i]].disp_B)/2, Vadd_disp/2);} //corrected for <=0
			new_ind[nn].C1 = (cur_ind[i].C1+cur_ind[MI[i]].C1)/2; //IF THIS IS GOING TO BE HERITABLE, WE WILL HAVE TO FIX THIS TOO
			new_ind[nn].C2 = (cur_ind[i].C2+cur_ind[MI[i]].C2)/2;

			//mutations(new_ind, nn, mut_rate, xDim);
	 		nn++;//increment new_ind counter
		}
   	}
 }
 new_n=nn;//nn less than new_n because some cells have no mates
 Free(n_offs);
 Free(MI);
 new_to_curr(new_n);
 return(new_n);

}

void mutations(ind * individuals, int id, double mut_rate, int xDim)
{
 double mut;

//Mutate the 1st allele for choosiness

// mut = Rf_rbinom(1,mut_rate);

// if (mut == 1){

//	 do{

//		individuals[id].C1 = individuals[id].C1 + Rf_runif(-0.1, 0.1);

//	 }while (individuals[id].C1 < 0.0 || individuals[id].C1 > 1.0 );

// }

   //Mutate the 2nd allele for choosiness

// mut = Rf_rbinom(1,mut_rate);

// if (mut == 1){

//	 do{

//		individuals[id].C2 = individuals[id].C2 + Rf_runif(-0.1, 0.1);

//	 }while (individuals[id].C2 < 0.0 || individuals[id].C2 > 1.0 );

// }

}


//mort = dispersal mortality
void disperse(int xDim, double mort, int popsize)
{
 double dist, meanDist, angle, new_x;

 int n = popsize;
 //Every individual moves.
 //Loop through each individual
 for (int i = 0; i < n; i++){
	 //Calculate destination
	 meanDist = cur_ind[i].disp_B;
	 dist = Rf_rnorm(0, meanDist);
	 cur_ind[i].alive = disp_mort(dist);
	 if (cur_ind[i].alive==0) { continue; }
	 new_x=dist+cur_ind[i].x_coord;
	 //Reflective boundaries on the x
	 if (new_x<0) { new_x=-1*new_x; }
	 while (new_x>xDim) { new_x=2*xDim-new_x; }
	 //printf("oldx=%f newx=%f oldy=%f newy=%f \n", cur_ind[i].x_coord, new_x, cur_ind[i].y_coord, new_y);
	 cur_ind[i].x_coord = new_x;
	 if (new_x > xDim || new_x<0) {
		 printf("error in x\n");
		 printf("%f %f\n", new_x, cur_ind[i].x_coord);
	 }

 }
	grid_update(n, xDim);
}


/* quick function to print out the individuals */

void check_inds (ind *individuals, int spX, int n) {
	int ii; //some iteration variables
	for (ii = 0; ii < n; ii++) {
		printf("individual %i -- X = %f -- HI = %f dispB = %f",ii,individuals[ii].x_coord,
			individuals[ii].HI, individuals[ii].disp_B);
		printf("\n\n");
	}

} 

void check_grid(int ** gd, int spX, int n){
	int xx;
	for (xx=0; xx<spX; xx++){
			printf("Grid cell %i, N=%i \n", xx, gd[xx]);
	}
}

void check_indices(int ** gd, int *** in, int spX){
	int x, z;
	for (x=0; x<spX; x++){
			printf("Grid cell %i\n", x);
			for (z=0; z<gd[x]; z++){
				printf("%i ", in[x][z]);
				if (z==gd[x]-1){printf("\n\n");}
			}
	}
}

/* function to return the first empty postion in an int vector */
int first_empty(int *tt, int len){
	for (int ii=0; ii<len; ii++) {
		if (tt[ii] == -99) { return(ii); }
	}
	return(-99);
}

//copy the new individuals over the current 
void new_to_curr(int nn) {
	cur_ind = Realloc(cur_ind, nn, ind);
	for (int ii = 0; ii < nn; ii++) {
		cur_ind[ii].x_coord = new_ind[ii].x_coord; cur_ind[ii].alive = 1;
		cur_ind[ii].HI=new_ind[ii].HI; cur_ind[ii].disp_B=new_ind[ii].disp_B;
		cur_ind[ii].C1=new_ind[ii].C1; cur_ind[ii].C2=new_ind[ii].C2;
	}
	new_ind = Realloc(new_ind, 1, ind); //move pointer
}

void output_individuals(int n, char mname[10], SEXP env) {
	double *out; SEXP ans;
	int ncol = 6; //define the number of columns
	PROTECT(ans = allocMatrix(REALSXP, n, ncol));
	out = REAL(ans); //pointer to output dataset
	for (int ii=0; ii<n; ii++) { //cycle through each of the individuals
		if (cur_ind[ii].alive != 0){ //if the individual is alive, copy the necessary info
			out[ii + 0 * n] = cur_ind[ii].x_coord; // set the xy values
			out[ii+1*n] = cur_ind[ii].alive; out[ii+2*n] = cur_ind[ii].HI;
			out[ii+3*n] = cur_ind[ii].disp_B;
			out[ii+4*n] = cur_ind[ii].C1; out[ii+5*n] = cur_ind[ii].C2;
		} else { // if ind dead.. set everything to -9
			for (int jj=0; jj<ncol; jj++) { out[ii + jj * n] = -9; }
		}
	}
	defineVar(install(mname), ans, env);	
	UNPROTECT(1);
}


// Provides a random draw from the collapsed 2D scalar Gaussian
//double c_gauss_rng(double sd){
//	double rn, x;
//	rn=Rf_runif(0,1);
//	x=sd*sqrt(2*log(1/(1-rn)));
//	return(x);
//}

// calculates probabilistic survival given a dispersal distance
int disp_mort(double dist){
	double p;
	double B=5;
	double A=1;
	p=1/(1+exp(-(dist-B)*A)); //mortality probability
	return(Rf_rbinom(1, 1-p)); //Alive or dead
}


// Summarises dispersal by x column and outputs to R
void disp_summ(char vname[10], int popsize, int xDim, SEXP env){
	double* dcol;
	SEXP ans;
	PROTECT(ans=allocMatrix(REALSXP, xDim, 1));
	dcol=REAL(ans);
	for (int xx=0; xx<xDim; xx++) {
		dcol[xx]=0;
	}
	for (int jj=0; jj<popsize; jj++) {
		if (cur_ind[jj].alive==1) {
			dcol[(int)cur_ind[jj].x_coord] += cur_ind[jj].disp_B;
		}
	}
	for (int xx=0; xx<xDim; xx++) {
		int n=0;
        n += density[xx];
		dcol[xx]=dcol[xx]/n; //mean dispersal by column
	}
	defineVar(install(vname), ans, env);
	UNPROTECT(1);
}


/* define the mother function
	R_n is the starting number of individuals
	R_spX & R_spY are the dimensions of the landscape
	R_d_cost is the dispersal cost */

SEXP mother(SEXP R_num_gens, SEXP R_n, SEXP R_spX, SEXP R_d_cost, SEXP R_lambda, SEXP R_K, SEXP R_m_rate,
			SEXP R_maxcost, SEXP R_beta, SEXP rho) {
	GetRNGstate();
	int gen = INTEGER(coerceVector(R_num_gens,INTSXP))[0]; //number of generations
	int n = INTEGER(coerceVector(R_n,INTSXP))[0]; //size of the habitat matrix - X
	int spX = INTEGER(coerceVector(R_spX,INTSXP))[0]; //size of the habitat matrix - X
	int lambda = INTEGER(coerceVector(R_lambda,INTSXP))[0]; //lambda
	int K = INTEGER(coerceVector(R_K,INTSXP))[0]; //K
	double d_cost = NUMERIC_VALUE(R_d_cost);//dispersal cost
	double m_rate = NUMERIC_VALUE(R_m_rate);
	double maxcost = NUMERIC_VALUE(R_maxcost);
	double beta = NUMERIC_VALUE(R_beta);
	double Vadd_disp=2; //Additive variance in dispersal
	
	printf("gen n spX lambda K d_cost m_rate maxcost alpha\n%i %i %i %i %i %f %f %f %f\n", gen, n, spX, lambda, K, d_cost, m_rate, maxcost, alpha);
	

	initialise(n, spX/2, spX); //also initialises the density grid
	
	
	for (int i=0; i < gen; i++) {
		n=reproduce(maxcost, beta, m_rate, spX, n,lambda, K, Vadd_disp);
		disperse(spX, d_cost, n);
		
		
		
				
		printf("n=%i  i=%i\n",n, i);
	}
	//check_grid(density, spX, spY, n);
	//check_inds(cur_ind, spX, spY, n);
	//check_indices(density, indices, spX, spY);
	
	output_individuals(n, "inds", rho);
	disp_summ("disp", n, spX, rho);

	

	//return some useless info
	SEXP res; PROTECT(res = allocVector(INTSXP,1));
	INTEGER(res)[0] = -9999;

	PutRNGstate();
	//release and free up memory 
	
	Free(cur_ind);
	Free(new_ind);
	for (int xx=0; xx<spX; xx++) {Free(density[xx]);} Free(density);
	for (int xx=0; xx<spX; xx++) {Free(indices[xx]);} Free(indices);

	UNPROTECT(1);
	return res;

}



